package org.huangrui.spark.java.core.rdd.dep;

import org.apache.spark.Partitioner;
import org.apache.spark.SparkConf;
import org.apache.spark.api.java.JavaPairRDD;
import org.apache.spark.api.java.JavaRDD;
import org.apache.spark.api.java.JavaSparkContext;
import scala.Tuple2;

import java.util.ArrayList;
import java.util.List;
import java.util.Objects;

/**
 * @Author hr
 * @Create 2024-10-19 16:50
 */
public class Spark08_Part_1 {
    public static void main(String[] args) throws InterruptedException {
        final SparkConf conf = new SparkConf();
        conf.setMaster("local[*]");
        conf.setAppName("spark");

        final JavaSparkContext jsc = new JavaSparkContext(conf);

        final List<Tuple2<String, Integer>> datas = new ArrayList<>();
        datas.add(new Tuple2<String, Integer>("nba", 1));
        datas.add(new Tuple2<String, Integer>("cba", 2));
        datas.add(new Tuple2<String, Integer>("nba", 3));
        datas.add(new Tuple2<String, Integer>("wnba", 4));

        final JavaRDD<Tuple2<String, Integer>> rdd = jsc.parallelize(datas, 2);
        final JavaPairRDD<String, Integer> mapRdd = rdd.mapToPair(t -> {
            System.out.println("==============================");
            return t;
        });

        // final JavaPairRDD<String, Integer> reduceRdd = mapRdd.reduceByKey(Integer::sum);//mapRDD.reduceByKey(new MyPartitioner(3), Integer::sum);
        // final JavaPairRDD<String, Integer> reduceRdd1 = reduceRdd.reduceByKey(Integer::sum);//reduceRDD.reduceByKey(new MyPartitioner(3), Integer::sum);
        final JavaPairRDD<String, Integer> reduceRdd = mapRdd.reduceByKey(new MyPartitioner(3), Integer::sum);
        final JavaPairRDD<String, Integer> reduceRdd1 = reduceRdd.reduceByKey(new MyPartitioner(3), Integer::sum);
        reduceRdd1.collect();

        System.out.println("计算完毕");
        // Stage = 1 + 1 = 2
        // Task  = 2 + 3 = 5
        // http://localhost:4040
        Thread.sleep(100000000L);

        jsc.close();
    }
}
// TODO 自定义分区器
//     1. 创建自定义类
//     2. 继承抽象类 Partitioner
//     3. 重写方法（2 + 2）
//         Partitioner(2) + Object(2)
//     4. 构建对象，在算子中使用
class MyPartitioner extends Partitioner {
    private final int numPartitions;
    public MyPartitioner(int numPartitions) {
        this.numPartitions = numPartitions;
    }
    // TODO 指定分区的数量
    @Override
    public int numPartitions() {
        return this.numPartitions;
    }

    // TODO 根据数据的KEY来获取数据存储的分区编号，编号从0开始
    @Override
    public int getPartition(Object key) {
        if ("nba".equals(key)) {
            return 1;
        }
        if ("cba".equals(key) || "wnba".equals(key)) {
            return 2;
        }
        return 0;
    }

    @Override
    public boolean equals(Object o) {
        if (this == o) {
            return true;
        }
        if (o == null || getClass() != o.getClass()) {
            return false;
        }
        MyPartitioner that = (MyPartitioner) o;
        return numPartitions == that.numPartitions;
    }

    @Override
    public int hashCode() {
        return Objects.hash(numPartitions);
    }
}
